import numpy as np
import torch
import torch.nn as nn
from torch import optim

from transformers.models.gpt2.modeling_gpt2 import GPT2Model
from transformers import BertTokenizer, BertModel
from einops import rearrange
# from embed import DataEmbedding, DataEmbedding_wo_time
from transformers.models.gpt2.configuration_gpt2 import GPT2Config


# 单例模式加载预训练模型

def singleton(cls):
    _instance = {}
    def _singleton(*args,**kwargs):
        #先判断这个类有没有对象
        if cls not in _instance:
            _instance[cls] = cls(*args, **kwargs)
        return _instance[cls]
    return _singleton()

@singleton
class LoadGPTModel:
    def __new__(cls, *args, **kwargs):
        gpt = GPT2Model.from_pretrained('gpt2', output_attentions=True, output_hidden_states=True)
        return gpt

class GPT4TS(nn.Module):
    
    def __init__(self, device):
        super(GPT4TS, self).__init__()
        self.is_gpt = 1                    # type=int, default=1
        self.patch_size = 2            # type=int, default=16
        self.pretrain = 1                # type=int, default=1
        self.stride = 8                    # type=int, default=8
        self.gpt_layers = 3                    # type=int, default=3
        self.d_model = 768
        self.seq_len = 96
        self.label_len = 48
        self.pred_len = 24
        self.freeze = 1

        self.patch_num = (self.seq_len - self.patch_size) // self.stride + 1

        self.padding_patch_layer = nn.ReplicationPad1d((0, self.stride)) 
        self.patch_num += 1
        # self.mlp = configs.mlp
        
        if self.is_gpt:
            if self.pretrain:
                # 加载gpt2-base预训练模型
                self.gpt2 = GPT2Model.from_pretrained('gpt2', output_attentions=True, output_hidden_states=True)
                # self.gpt2 = LoadGPTModel()
            else:
                print("------------------no pretrain------------------")
                self.gpt2 = GPT2Model(GPT2Config())

            self.gpt2.h = self.gpt2.h[:self.gpt_layers]
            print("gpt2 = {}".format(self.gpt2))
        
        self.in_layer = nn.Linear(self.patch_size, self.d_model)
        self.out_layer = nn.Linear(self.d_model * self.patch_num, self.pred_len)
        
        if self.freeze and self.pretrain:
            for i, (name, param) in enumerate(self.gpt2.named_parameters()):
                if 'ln' in name or 'wpe' in name:
                    param.requires_grad = True
                else:
                    param.requires_grad = False

        for layer in (self.gpt2, self.in_layer, self.out_layer):
            layer.to(device=device)
            layer.train()
        
        self.cnt = 0

    def forward(self, x, exogenous_data=None, device="cpu", y_hist=None):

        x = x.reshape(x.size(0), x.size(1), x.size(2))

        B, L, M = x.shape

        means = x.mean(1, keepdim=True).detach()
        x = x - means
        stdev = torch.sqrt(torch.var(x, dim=1, keepdim=True, unbiased=False)+ 1e-5).detach() 
        x /= stdev

        x = rearrange(x, 'b l m -> b m l')

        x = self.padding_patch_layer(x)
        x = x.unfold(dimension=-1, size=self.patch_size, step=self.stride)
        x = rearrange(x, 'b m n p -> (b m) n p')

        outputs = self.in_layer(x)
        if self.is_gpt:
            outputs = self.gpt2(inputs_embeds=outputs).last_hidden_state

        outputs = self.out_layer(outputs.reshape(B*M, -1))
        outputs = rearrange(outputs, '(b m) l -> b l m', b=B)

        outputs = outputs * stdev
        outputs = outputs + means

        return outputs
